[FEATURE] support Wan2.2 LoRA finetune #1418
Conversation
zhtmike
changed the title
zhtmike
requested review from
CaitinZhao,
SamitHuang,
vigo999 and
zhanghuiyao
as code owners
|
/gemini review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for LoRA fine-tuning for the Wan2.2 5B model. It adds a new training script, dataset handling, and a LoRATrainer class that encapsulates the fine-tuning logic. The changes also include necessary modifications to the VAE and model components to support training and LoRA weight loading for inference.
My review focuses on the correctness and robustness of the new training pipeline. I've found a bug in the validation logic that could lead to incorrect behavior during training, and a potential critical issue in the core algorithm logic. I've also pointed out some typos in the new documentation. Overall, the changes are well-structured and provide a solid foundation for fine-tuning.
There was a problem hiding this comment.
Pull request overview
This PR adds LoRA (Low-Rank Adaptation) finetuning support for the Wan2.2 5B Text-to-Video model, enabling efficient parameter-efficient training on custom datasets.
Key Changes
- Implements
LoRATrainerclass with ZeRO-3 optimization support for distributed training - Adds
VideoDatasetfor loading video data from CSV or file-based formats - Integrates LoRA checkpoint loading into the generation pipeline
Reviewed changes
Copilot reviewed 15 out of 15 changed files in this pull request and generated 12 comments.
Show a summary per file
| File | Description |
|---|---|
| mindone/trainers/zero.py | Updates optimizer parameter tuple retrieval logic for MindSpore 2.7.0+ compatibility |
| examples/wan2_2/wan/utils/fm_solvers_unipc.py | Rounds timesteps to ensure integer values |
| examples/wan2_2/wan/utils/fm_solvers.py | Rounds timesteps to ensure integer values |
| examples/wan2_2/wan/trainer/utils.py | Adds utility functions for gradient clipping, checkpoint saving, and distributed operations |
| examples/wan2_2/wan/trainer/lora_trainer.py | Implements the main LoRA training loop with validation and checkpoint management |
| examples/wan2_2/wan/trainer/dataset.py | Implements video dataset loader with preprocessing and bucketing support |
| examples/wan2_2/wan/trainer/init.py | Exports trainer and dataset creation functions |
| examples/wan2_2/wan/textimage2video.py | Adds LoRA checkpoint loading capability to the pipeline |
| examples/wan2_2/wan/modules/vae2_2.py | Adds return_log_var parameter to VAE encode method |
| examples/wan2_2/wan/modules/vae2_1.py | Adds return_log_var parameter to VAE encode method |
| examples/wan2_2/wan/modules/model.py | Enables gradient checkpointing during training to reduce memory usage |
| examples/wan2_2/train.py | Implements training script with distributed training support |
| examples/wan2_2/scripts/train_lora_2p.sh | Adds shell script for launching 2-GPU LoRA training |
| examples/wan2_2/generate.py | Adds --lora_dir argument for loading LoRA weights during inference |
| examples/wan2_2/finetune.md | Adds comprehensive documentation for LoRA finetuning workflow |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
You can also share your feedback on Copilot code review for a chance to win a $100 gift card. Take the survey.
1 participant
Support Wan2.2 5B LoRA Finetune
Require PR #1383
What does this PR do?
Fixes # (issue)
Adds # (feature)
Before submitting
What's New. Here are thedocumentation guidelines
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@xxx